# -*- coding: utf-8 -*-
import random
import torch
import numpy as np
import torch as t
import torch.nn as nn
import sys
import torch.optim as optim


def get_data(pairs, paths_dict, p):
    # 记录随机索引下的路径
    paths_list = []
    # 记录pairs中的标签
    label_list = []
    for pair in pairs:
        if (pair[0], pair[1]) in paths_dict:
            # 判断该节点对是否有对应的路径
            if len(paths_dict[(pair[0], pair[1])]):
                # 将该节点对对应的路径赋值给变量 paths
                paths = paths_dict[(pair[0], pair[1])]
                # 从给定的路径列表 paths 中随机选择一些路径作为训练集，并将剩余路径作为验证集。
                # 如果 paths 中的路径数大于或等于 p，则从中随机选择 p 个路径作为训练集，剩余路径作为验证集；
                # 否则，从中随机选择 p 个路径作为训练集，可以重复选择路径，剩余路径作为验证集。
                if len(paths) >= p:
                    # 从路径列表paths中随机选择5条路径，并返回它们在路径列表中的索引，replace=False时，表示不允许重复选择路径
                    indices = np.random.choice(len(paths), p, replace=False)
                else:
                    # 从路径列表paths中随机选择5条路径，并返回它们在路径列表中的索引，replace=True时，表示允许重复选择路径
                    indices = np.random.choice(len(paths), p, replace=True)
                # 记录随机选择的5条路径，与上面的indices有关，对应索引
                paths_list.append([paths[i] for i in indices])
            else:
                # 直接记录空数据
                paths_list.append([])
        else:
            # 直接记录空数据
            paths_list.append([])

        label_list.append(pair[2])
    return paths_list, label_list
def recommend_attractions(district, purpose, type, season, age, transportation):

    params=district*2880+purpose*720+type*144+season*24+age*4+transportation;
    return params


# def eval_ctr1(model, pairs, paths_dict, p, batch_size):
#     model.eval()
#
#     pred_label = []
#     paths_list = get_data(pairs, paths_dict, p)[0]
#
#     for i in range(0, len(pairs), batch_size):
#         batch_label = model(paths_list[i: i+batch_size]).cpu().detach().numpy().tolist()
#         pred_label.extend(batch_label)
#
#
#     # 将预测值替换回评分列
#     predictions = np.array(pred_label)
#     test = np.array(pairs)
#     test[:, 2] = predictions
#
#     # 根据预测值从大到小对数据排序
#     sorted_test = test[test[:, 2].argsort()[::-1]]
#
#
#     # second_column_values = sorted_test[:, 1]
#     # result1 = second_column_values[:5]
#
#     result1 = []
#
#     # 遍历排序后的数据，将满足条件的第二列数据放入结果数组中
#     for data in sorted_test:
#         if len(result1) >= 5:
#             break
#         result1.append(data[1])
#
#     # 如果满足条件的数据多于数组剩余容量的大小，随机选择数据放入结果数组中
#     if len(result1) > 5:
#         result1 = random.sample(result1, 5)
#
#
#     # result1 = random.choices([data[1] for data in test], k=5)
#     print(','.join(map(str, result1)))
#     return result1

def eval_ctr1(model, pairs, paths_dict, p, batch_size):
    model.eval()

    pred_label = []
    paths_list = get_data(pairs, paths_dict, p)[0]

    for i in range(0, len(pairs), batch_size):
        batch_label = model(paths_list[i: i+batch_size]).cpu().detach().numpy().tolist()
        pred_label.extend(batch_label)


    # 将预测值替换回评分列
    predictions = np.array(pred_label)
    # 添加代码
    # sorted_predictions = sorted(predictions, reverse=True)
    # print(','.join(map(str, sorted_predictions)))
    test = np.array(pairs)
    test[:, 2] = predictions

    # 根据预测值从大到小对数据排序
    sorted_test = test[test[:, 2].argsort()[::-1]]
    # print(','.join(map(str, sorted_test)))

    # for row in sorted_test:
    #     col1 = row[0]  # 获取第一列的值
    #     col2 = row[1]  # 获取第二列的值
    #     col3 = row[2]  # 获取第三列的值
    #
    #     print(col1, col2, col3)
    #
    # second_column_values = sorted_test[:, 1]
    # result1 = second_column_values[:5]
    # print(result1)
    result1 = []

    # result1 = random.choices([data[1] for data in test], k=5)


    for data in sorted_test:
        if len(result1) >= 5:
            break
        result1.append(data[1])

    # 如果满足条件的数据多于数组剩余容量的大小，随机选择数据放入结果数组中
    if len(result1) > 5:
        result1 = random.sample(result1, 5)

    # print(result1)


    return result1

class RKGE(nn.Module):

    def __init__(self, n_entity, dim):
        super(RKGE, self).__init__()
        self.dim = dim
        self.entity_embedding_matrix = nn.Parameter(t.randn(n_entity, dim))
        self.rnn = nn.RNN(dim, dim, num_layers=1, bidirectional=True)
        # 注意力权重
        self.M = nn.Linear(2*dim, 1)
        # 将大小为dim的输入向量转换为大小为dim的输出向量
        self.W = nn.Linear(dim, dim, bias=False)

        self.H = nn.Linear(dim, dim)
        self.Wr = nn.Linear(dim, 1)
    # ???
    def forward(self, paths_list):

        path_embeddings = self.get_path_embedding(paths_list)

        # (len, -1, 2 * dim)
        hidden_states = self.rnn(path_embeddings)[0]

        for i in range(4):
            attention = t.sigmoid(self.M(hidden_states[0]))

            if i == 0:
                # sigmoid 函数用于将线性变换的结果映射到一个注意力分数,范围为 [0, 1]
                candidate_hidden_states = t.sigmoid(self.H(path_embeddings[i]))
                next_hidden_states = attention * candidate_hidden_states
            else:
                candidate_hidden_states = t.sigmoid(self.W(next_hidden_states) + self.H(path_embeddings[i]))
                next_hidden_states = (1 - attention) * next_hidden_states + attention * candidate_hidden_states
        aggregation_hidden_states = next_hidden_states.reshape(len(paths_list), -1, self.dim).mean(dim=1)
        predicts = t.sigmoid(self.Wr(aggregation_hidden_states)).reshape(-1)

        return predicts
    # 用到了这里的函数
    def get_path_embedding(self, paths_list):
        path_embedding_list = []
        # 创建一个形状为(4, 5, dim(20))的张量
        zeros = t.zeros(4, 5, self.dim)
        if t.cuda.is_available():
            zeros = zeros.to(self.entity_embedding_matrix.data.device)
        for paths in paths_list:
            # 4*20的张量，每一行表示一个节点的嵌入向量。
            # 节点的值相同时，获得的嵌入向量值也相同
            entity_embedding_list = []
            for path in paths:
                entity_embedding_list.append(self.entity_embedding_matrix[path].reshape(-1, 1, self.dim))

            # (len, n_path, dim)
            if len(paths) == 0:
                path_embedding_list.append(zeros)
                continue
            path_embedding_list.append(t.cat(entity_embedding_list, dim=1))
        # t.cat用于将多个张量（tensor）沿着指定的维度(dim=1)进行拼接
        return t.cat(path_embedding_list, dim=1)

class Configuration:
    def __init__(self):
        self.dataset = 'zhishitupu'
        self.ratio = 0.2
        self.lr = 5e-3
        self.l2 = 1e-4
        self.batch_size = 1024
        self.device = 'cuda:0'
        self.dim = 20
        self.p = 5


# 创建一个配置对象
config = Configuration()

if __name__ == "__main__":
    district=1
    purpose=0
    type=1
    season=2
    age=2
    transportation=3

    # district =int(sys.argv[1])
    # purpose = int(sys.argv[2])
    # type=int(sys.argv[3])
    # season=int(sys.argv[4])
    # age=int(sys.argv[5])
    # transportation=int(sys.argv[6])
    # district=1
    # purpose=1
    # type=1
    # season=1
    # age=1
    # transportation=1



    # data_dir = './RKGE-master/data/' + config.dataset + '/'
    # zfc = recommend_attractions(district, purpose, type, season, age, transportation)
    # paths_dict = np.load(data_dir + str(config.ratio) + '_3_path_dict.npy', allow_pickle=True).item()
    # entity_list = np.load(data_dir + '_entity_list.npy').tolist()
    # n_entity = len(entity_list)
    # test = np.load(data_dir + str(config.ratio) + '_test.npy').tolist()
    # device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    # # 加载模型参数（使用相对路径或绝对路径）
    # # state_dict = torch.load('./RKGE-master/model.pth', map_location=device)
    # # model = RKGE(n_entity, config.dim)
    # # model.load_state_dict(state_dict)
    # # ??????
    # state_dict = torch.load('./RKGE-master/model.pth', map_location=device)
    # model = RKGE(n_entity, config.dim)
    #
    # # # 查看检查点中的参数形状
    # # checkpoint_shapes = {k: v.shape for k, v in state_dict.items()}
    # #
    # # # 检查 "entity_embedding_matrix" 参数是否存在
    # # if 'entity_embedding_matrix' in checkpoint_shapes:
    # #     # 获取当前模型中 "entity_embedding_matrix" 的形状
    # #     current_shape = model.entity_embedding_matrix.shape
    # #
    # #     # 检查形状是否匹配
    # #     if current_shape != checkpoint_shapes['entity_embedding_matrix']:
    # #         # 如果形状不匹配，调整当前模型的 "entity_embedding_matrix" 形状
    # #         model.entity_embedding_matrix = torch.nn.Parameter(
    # #             state_dict['entity_embedding_matrix'].data[:current_shape[0]]
    # #         )
    #
    # # 加载其他参数
    # model.load_state_dict(state_dict)
    #
    #
    #
    #
    # for row in test:
    #     row[0] = zfc
    #
    # sort_data = eval_ctr1(model, test, paths_dict, config.p, config.batch_size)
    # recommended_ids=sort_data
    # print(','.join(map(str, recommended_ids)))

    # data_dir = './RKGE-master/data/' + config.dataset + '/'
    data_dir = './data/' + config.dataset + '/'
    zfc = recommend_attractions(district, purpose, type, season, age, transportation)
    paths_dict = np.load(data_dir + str(config.ratio) + '_3_path_dict.npy', allow_pickle=True).item()
    # paths_dict1 = np.load(data_dir + '1' + '_3_path_dict.npy', allow_pickle=True).item()
    # 遍历原始数据
    for key, value in paths_dict.items():
        # 检查是否已存在相同键的数据
        if key not in paths_dict:
            paths_dict[key] = value
    entity_list = np.load(data_dir + '_entity_list.npy').tolist()
    n_entity = len(entity_list)
    test = np.load(data_dir + str(config.ratio) + '_test.npy').tolist()
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    # state_dict = torch.load('./RKGE-master/model.pth', map_location=device)
    state_dict = torch.load('model.pth', map_location=device)
    model = RKGE(n_entity, config.dim)
    model.load_state_dict(state_dict)

    for row in test:
        row[0] = 3758

    sort_data = eval_ctr1(model, test, paths_dict, config.p, config.batch_size)
    recommended_ids = sort_data

    # if district==1:
    #     district_values = []
    #     # 读取 kg.txt 文件
    #     with open(data_dir + 'kg.txt', 'r') as file:
    #         for line in file:
    #             # 分割每一行的数据
    #             data = line.strip().split('\t')
    #             if len(data) == 3 and data[1] == '177':
    #                 # 找到第二列值为 177 的对应的第一列数据，并存入数组
    #                 district_values.append(data[0])
    #     while len(recommended_ids) < 5:
    #         # 从 district_values 随机选择一个元素
    #         random_id = random.choice(district_values)
    #         if random_id not in recommended_ids:
    #             recommended_ids.append(random_id)
    #
    # elif district==2:
    #     district_values = []
    #     # 读取 kg.txt 文件
    #     with open(data_dir + 'kg.txt', 'r') as file:
    #         for line in file:
    #             # 分割每一行的数据
    #             data = line.strip().split('\t')
    #             if len(data) == 3 and data[1] == '176':
    #                 # 找到第二列值为 177 的对应的第一列数据，并存入数组
    #                 district_values.append(data[0])
    #     while len(recommended_ids) < 5:
    #         # 从 district_values 随机选择一个元素
    #         random_id = random.choice(district_values)
    #         if random_id not in recommended_ids:
    #             recommended_ids.append(random_id)
    #
    # elif district==3:
    #     district_values = []
    #     # 读取 kg.txt 文件
    #     with open(data_dir + 'kg.txt', 'r') as file:
    #         for line in file:
    #             # 分割每一行的数据
    #             data = line.strip().split('\t')
    #             if len(data) == 3 and data[1] == '170':
    #                 # 找到第二列值为 177 的对应的第一列数据，并存入数组
    #                 district_values.append(data[0])
    #     while len(recommended_ids) < 5:
    #         # 从 district_values 随机选择一个元素
    #         random_id = random.choice(district_values)
    #         if random_id not in recommended_ids:
    #             recommended_ids.append(random_id)
    #
    # elif district==4:
    #     district_values = []
    #     # 读取 kg.txt 文件
    #     with open(data_dir + 'kg.txt', 'r') as file:
    #         for line in file:
    #             # 分割每一行的数据
    #             data = line.strip().split('\t')
    #             if len(data) == 3 and data[1] == '169':
    #                 # 找到第二列值为 177 的对应的第一列数据，并存入数组
    #                 district_values.append(data[0])
    #     while len(recommended_ids) < 5:
    #         # 从 district_values 随机选择一个元素
    #         random_id = random.choice(district_values)
    #         if random_id not in recommended_ids:
    #             recommended_ids.append(random_id)
    #
    # elif district==5:
    #     district_values = []
    #     # 读取 kg.txt 文件
    #     with open(data_dir + 'kg.txt', 'r') as file:
    #         for line in file:
    #             # 分割每一行的数据
    #             data = line.strip().split('\t')
    #             if len(data) == 3 and data[1] == '174':
    #                 # 找到第二列值为 177 的对应的第一列数据，并存入数组
    #                 district_values.append(data[0])
    #     while len(recommended_ids) < 5:
    #         # 从 district_values 随机选择一个元素
    #         random_id = random.choice(district_values)
    #         if random_id not in recommended_ids:
    #             recommended_ids.append(random_id)
    #
    # elif district==6:
    #     district_values = []
    #     # 读取 kg.txt 文件
    #     with open(data_dir + 'kg.txt', 'r') as file:
    #         for line in file:
    #             # 分割每一行的数据
    #             data = line.strip().split('\t')
    #             if len(data) == 3 and data[1] == '175':
    #                 # 找到第二列值为 177 的对应的第一列数据，并存入数组
    #                 district_values.append(data[0])
    #     while len(recommended_ids) < 5:
    #         # 从 district_values 随机选择一个元素
    #         random_id = random.choice(district_values)
    #         if random_id not in recommended_ids:
    #             recommended_ids.append(random_id)
    #
    # elif district==7:
    #     district_values = []
    #     # 读取 kg.txt 文件
    #     with open(data_dir + 'kg.txt', 'r') as file:
    #         for line in file:
    #             # 分割每一行的数据
    #             data = line.strip().split('\t')
    #             if len(data) == 3 and data[1] == '178':
    #                 # 找到第二列值为 177 的对应的第一列数据，并存入数组
    #                 district_values.append(data[0])
    #     while len(recommended_ids) < 5:
    #         # 从 district_values 随机选择一个元素
    #         random_id = random.choice(district_values)
    #         if random_id not in recommended_ids:
    #             recommended_ids.append(random_id)
    #
    # elif district==8:
    #     district_values = []
    #     # 读取 kg.txt 文件
    #     with open(data_dir + 'kg.txt', 'r') as file:
    #         for line in file:
    #             # 分割每一行的数据
    #             data = line.strip().split('\t')
    #             if len(data) == 3 and data[1] == '171':
    #                 # 找到第二列值为 177 的对应的第一列数据，并存入数组
    #                 district_values.append(data[0])
    #     while len(recommended_ids) < 5:
    #         # 从 district_values 随机选择一个元素
    #         random_id = random.choice(district_values)
    #         if random_id not in recommended_ids:
    #             recommended_ids.append(random_id)
    #
    #
    # elif district==9:
    #     district_values = []
    #     # 读取 kg.txt 文件
    #     with open(data_dir + 'kg.txt', 'r') as file:
    #         for line in file:
    #             # 分割每一行的数据
    #             data = line.strip().split('\t')
    #             if len(data) == 3 and data[1] == '173':
    #                 # 找到第二列值为 177 的对应的第一列数据，并存入数组
    #                 district_values.append(data[0])
    #     while len(recommended_ids) < 5:
    #         # 从 district_values 随机选择一个元素
    #         random_id = random.choice(district_values)
    #         if random_id not in recommended_ids:
    #             recommended_ids.append(random_id)
    #
    # elif district==10:
    #     district_values = []
    #     # 读取 kg.txt 文件
    #     with open(data_dir + 'kg.txt', 'r') as file:
    #         for line in file:
    #             # 分割每一行的数据
    #             data = line.strip().split('\t')
    #             if len(data) == 3 and data[1] == '172':
    #                 # 找到第二列值为 177 的对应的第一列数据，并存入数组
    #                 district_values.append(data[0])
    #     while len(recommended_ids) < 5:
    #         # 从 district_values 随机选择一个元素
    #         random_id = random.choice(district_values)
    #         if random_id not in recommended_ids:
    #             recommended_ids.append(random_id)
    #
    # elif district==11:
    #     district_values = []
    #     # 读取 kg.txt 文件
    #     with open(data_dir + 'kg.txt', 'r') as file:
    #         for line in file:
    #             # 分割每一行的数据
    #             data = line.strip().split('\t')
    #             if len(data) == 3 and data[1] == '168':
    #                 # 找到第二列值为 177 的对应的第一列数据，并存入数组
    #                 district_values.append(data[0])
    #     while len(recommended_ids) < 5:
    #         # 从 district_values 随机选择一个元素
    #         random_id = random.choice(district_values)
    #         if random_id not in recommended_ids:
    #             recommended_ids.append(random_id)

    print(','.join(map(str, recommended_ids)))